#include "http_server.h"
#include "constants.h"
#include "logger.h"

#include "hnswlib_index.h"
#include <rapidjson/writer.h>
#include <rapidjson/stringbuffer.h>

HttpServer::HttpServer(const std::string&host,int port,VectorDatabase* vector_database_,RaftStuff* raft_stuff):host(host),port(port),vector_database(vector_database_),raft_stuff_(raft_stuff){
    server.Post(
        "/search",[this](const httplib::Request& req, httplib::Response& res){
        searchHandler(req,res);
    });

    server.Post(
        "/insert",[this](const httplib::Request& req, httplib::Response& res){
            insertHandler(req,res);
        }
    );
    server.Post(
        "/upsert",[this](const httplib::Request& req, httplib::Response& res){
            upsertHandler(req,res);
        }
    );
    server.Post(
        "/query",[this](const httplib::Request& req, httplib::Response& res){
            queryHandler(req,res);
        }
    );
    server.Post(
        "/admin/snapshot",[this](const httplib::Request& req, httplib::Response& res){
            snapshotHandler(req,res);
        }
    );
    server.Post(
        "/admin/setLeader",[this](const httplib::Request& req, httplib::Response& res){
            setLeaderHandler(req,res);
        }
    );
    server.Post(
        "/admin/addFollower",[this](const httplib::Request& req, httplib::Response& res){
            addFollowerHandler(req,res);
        }
    );
    server.Get(
        "/admin/listNode",[this](const httplib::Request& req, httplib::Response& res){
            listNodeHandler(req,res);
        }
    );
    server.Get(
        "/admin/getNode",[this](const httplib::Request& req, httplib::Response& res){
            getNodeHandler(req,res);
        }
    );
}

void HttpServer::start(){
    server.listen(host.c_str(),port);
}

bool HttpServer::isRequestValid(const rapidjson::Document& json_request, CheckType check_type) {
    switch (check_type) {
        case CheckType::SEARCH:
            return json_request.HasMember(REQUEST_VECTORS) &&
                   json_request.HasMember(REQUEST_K) &&
                   (!json_request.HasMember(REQUEST_INDEX_TYPE) || json_request[REQUEST_INDEX_TYPE].IsString());
        case CheckType::INSERT:
            return json_request.HasMember(REQUEST_VECTORS) &&
                   json_request.HasMember(REQUEST_ID) &&
                   (!json_request.HasMember(REQUEST_INDEX_TYPE) || json_request[REQUEST_INDEX_TYPE].IsString());
        case CheckType::UPSERT:
            return json_request.HasMember(REQUEST_VECTORS) &&
                   json_request.HasMember(REQUEST_ID) &&
                   (!json_request.HasMember(REQUEST_INDEX_TYPE) || json_request[REQUEST_INDEX_TYPE].IsString());
        case CheckType::QUERY:
            return true;
        default:
            return false;
    }
}

IndexFactory::IndexType HttpServer::getIndexTypeFromRequest(const rapidjson::Document &req){
    // 获取索引类型
    if(req.HasMember(REQUEST_INDEX_TYPE)){
        const std::string index_type=req[REQUEST_INDEX_TYPE].GetString();
        if(index_type==INDEX_TYPE_FLAT){
            return IndexFactory::IndexType::FLAT;
        }else if(index_type==INDEX_TYPE_HNSW){
            return IndexFactory::IndexType::HNSW;
        }
    }
    return IndexFactory::IndexType::UNKNOWN;
}

void HttpServer::insertHandler(const httplib::Request &req, httplib::Response &res){
    GlobalLogger->debug("insertHandler");

    //解析请求
    rapidjson::Document json_req;
    json_req.Parse(req.body.c_str());

    GlobalLogger->info("insert request body:{}",req.body);
    
    //检查是否解析有效
    if(!json_req.IsObject()){
        GlobalLogger->error("insert request body is not a json object");
        res.status = 400;

        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"insert request body is not a json object");
        return;
    }

    //检查请求合法性
    if(!isRequestValid(json_req,CheckType::INSERT)){
        GlobalLogger->error("insert request body is not valid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"insert request body is not valid");
        return;
    }

    //获取插入参数
    std::vector<float> data;
    for(const auto &item:json_req[REQUEST_VECTORS].GetArray()){
        data.push_back(item.GetFloat());
    }

    uint64_t label=json_req[REQUEST_ID].GetUint64();

    GlobalLogger->debug("insert request: label={}",label);

    IndexFactory::IndexType index_type=getIndexTypeFromRequest(json_req);
    //索引类型为UNKNOWN，错误
    if(index_type==IndexFactory::IndexType::UNKNOWN){
        GlobalLogger->error("insert request: index type is not valid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"index type is not valid");
        return;
    }

    void* index=IndexFactory::getGlobalIndexFactory().getIndex(index_type);
    if(index==nullptr){
        GlobalLogger->error("index null");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"index null");
        return;
    }
    //根据索引类型进行不同的插入
    switch (index_type)
    {
    case IndexFactory::IndexType::FLAT:{
        FaissIndex* faissIndex=static_cast<FaissIndex*>(index);
        faissIndex->insert(data,label);
        break;
    }   
    case IndexFactory::IndexType::HNSW:{
        GlobalLogger->info("insert hnsw index");
        HNSWLibIndex* hnswLibIndex=static_cast<HNSWLibIndex*>(index);
        try {
            hnswLibIndex->insert_vectors(data, label);
        } catch (const std::exception& e) {
            GlobalLogger->error("HNSW insert failed: {}", e.what());
            res.status = 500;
            setErrorJsonResponse(res, RESPONSE_RETCODE_ERROR, "Insert failed: " + std::string(e.what()));
            return;
        }
        hnswLibIndex->insert_vectors(data,label);
        break;
    }
    default:
        break;
    }

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType &allocator = json_response.GetAllocator();

    //添加retcode
    json_response.AddMember(RESPONSE_RETCODE, RESPONSE_RETCODE_SUCCESS, allocator);
    setJsonResponse(res,json_response);
    GlobalLogger->info("insert success");
}   

void HttpServer::searchHandler(const httplib::Request &req, httplib::Response &res){
    GlobalLogger->debug("searchHandler start");
    rapidjson::Document json_req;
    json_req.Parse(req.body.c_str());
    GlobalLogger->info("searchHandler req:{}",req.body);
    //解析是否有效
    if(!json_req.IsObject()){
        res.status=400;
        GlobalLogger->error("json_req is not object");
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"json_req is not object");
        return;
    }

    //检查请求合法性
    if(!isRequestValid(json_req,CheckType::SEARCH)){
        GlobalLogger->error("request is not valid");
        res.status=400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"request is not valid");
        return;
    }

    std::vector<float> query;
    for(const auto& item:json_req[REQUEST_VECTORS].GetArray()){
        query.push_back(item.GetFloat());
    }

    int k=json_req[REQUEST_K].GetInt();
    GlobalLogger->info("query k:{}",k);

    IndexFactory::IndexType index_type=getIndexTypeFromRequest(json_req);
    if(index_type==IndexFactory::IndexType::UNKNOWN){
        GlobalLogger->error("unknown index type");
        res.status=400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"unknown index type");
        return;
    }

    std::pair<std::vector<long>,std::vector<float>> results=vector_database->search(json_req);

    //转换结果为JSON
    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    //检查结果是否有效
    bool valid_result = true;
    rapidjson::Value vectors(rapidjson::kArrayType);
    rapidjson::Value distances(rapidjson::kArrayType);

    for(size_t i=0;i<results.first.size();i++){
        if(results.first[i]!=-1){
            valid_result = true;
            vectors.PushBack(results.first[i],allocator);
            distances.PushBack(results.second[i],allocator);
        }
    }

    if(valid_result){
        json_response.AddMember(RESPONSE_VECTORS,vectors,allocator);
        json_response.AddMember(RESPONSE_DISTANCES,distances,allocator);
    }

    //设置响应
    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::setJsonResponse(httplib::Response &res,const rapidjson::Document &json_response){
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

    json_response.Accept(writer);
    res.set_content(buffer.GetString(),RESPONSE_CONTENT_TYPE_JSON);
}

void HttpServer::setErrorJsonResponse(httplib::Response &res,int errCode,const std::string &error_message){
    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    json_response.AddMember(RESPONSE_RETCODE,errCode,allocator);
    json_response.AddMember(RESPONSE_ERROR_MSG,rapidjson::StringRef(error_message.c_str()),allocator);
    setJsonResponse(res,json_response);
}
void HttpServer::upsertHandler(const httplib::Request &req, httplib::Response &res){
    GlobalLogger->debug("Receive upsert request");
    //解析相关请求
    rapidjson::Document json_request;
    json_request.Parse(req.body.c_str());

    //检查解析结果
    if(!json_request.IsObject()){
        GlobalLogger->debug("json request invalid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"json request invalid");
        return;
    }
    //检查请求合法性
    if(!isRequestValid(json_request,CheckType::UPSERT)){
        GlobalLogger->debug("Missing vectors or id parameter in the request");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"Missing vectors or id parameter in the request");
        return;
    }

    uint64_t label=json_request[REQUEST_ID].GetUint64();
    //获取索引类型
    IndexFactory::IndexType index_type=getIndexTypeFromRequest(json_request);

    raft_stuff_->appendEntries(req.body);
    vector_database->upsert_vectors(label,json_request,index_type);

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::queryHandler(const httplib::Request &req, httplib::Response &res){
    GlobalLogger->debug("Receive query request");
    rapidjson::Document json_request;
    json_request.Parse(req.body.c_str());
    GlobalLogger->info("queryHandler req:{}",req.body);

    if(!json_request.IsObject()){
        GlobalLogger->debug("json request invalid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"json request invalid");
        return;

    }

    if(!isRequestValid(json_request,CheckType::QUERY)){ 
        GlobalLogger->debug("request is not valid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"request is not valid");
        return;
    }

    std::vector<float> query;
    uint64_t id=json_request[REQUEST_ID].GetInt64();
    rapidjson::Document json_data=vector_database->query(id);

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    if(!json_data.IsNull()){
        for(auto it=json_data.MemberBegin();it!=json_data.MemberEnd();it++){
            json_response.AddMember(it->name,it->value,allocator);
        }
    }

    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::snapshotHandler(const httplib::Request &req, httplib::Response &res)
{
    GlobalLogger->debug("Receive snapshot request");
    vector_database->takeSnapshot();

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::setLeaderHandler(const httplib::Request &req, httplib::Response &res)
{
    GlobalLogger->debug("Receive set leader request");
    raft_stuff_->enableElectionTimeout(10000,20000);

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::addFollowerHandler(const httplib::Request &req, httplib::Response &res)
{
    GlobalLogger->debug("Receive add follower request");

    rapidjson::Document json_request;
    json_request.Parse(req.body.c_str());

    if(!json_request.IsObject()){
        GlobalLogger->debug("json request invalid");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"json request invalid");
        return;
    }

    if(!raft_stuff_->is_leader()){
        GlobalLogger->debug("Not leader");
        res.status = 400;
        setErrorJsonResponse(res,RESPONSE_RETCODE_ERROR,"Not leader");
        return;
    }

    int node_id=json_request["nodeId"].GetInt();
    std::string endpoint=json_request["endpoint"].GetString();

    raft_stuff_->addSrv(node_id,endpoint);

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();
    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::listNodeHandler(const httplib::Request &req, httplib::Response &res)
{
    GlobalLogger->debug("Receive list node request");
    auto nodes_info=raft_stuff_->getAllNodesInfo();

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    rapidjson::Value nodes_array(rapidjson::kArrayType);
    for(const auto& node_info:nodes_info){
        rapidjson::Value node_info_json(rapidjson::kObjectType);
        node_info_json.AddMember("nodeId",std::get<0>(node_info),allocator);
        node_info_json.AddMember("endpoint",rapidjson::Value(std::get<1>(node_info).c_str(),allocator),allocator);
        node_info_json.AddMember("status",rapidjson::Value(std::get<2>(node_info).c_str(),allocator),allocator);
        node_info_json.AddMember("lastLogIndex",std::get<3>(node_info),allocator);
        node_info_json.AddMember("lastSuccResp",std::get<4>(node_info),allocator);
        nodes_array.PushBack(node_info_json,allocator);
    }
    json_response.AddMember("nodes",nodes_array,allocator);

    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}

void HttpServer::getNodeHandler(const httplib::Request &req, httplib::Response &res)
{
    GlobalLogger->debug("Receive get node request");

    std::tuple<int,std::string,std::string,nuraft::ulong,nuraft::ulong> node_info=raft_stuff_->getCurrentNodesInfo();

    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    rapidjson::Value nodes_array(rapidjson::kArrayType);
    rapidjson::Value node_object(rapidjson::kObjectType);
    node_object.AddMember("nodeId",std::get<0>(node_info),allocator);
    node_object.AddMember("endpoint",rapidjson::Value(std::get<1>(node_info).c_str(),allocator),allocator);
    node_object.AddMember("status",rapidjson::Value(std::get<2>(node_info).c_str(),allocator),allocator);
    node_object.AddMember("lastLogIndex",std::get<3>(node_info),allocator);
    node_object.AddMember("lastSuccResp",std::get<4>(node_info),allocator);

    json_response.AddMember("node",node_object,allocator);
    json_response.AddMember(RESPONSE_RETCODE,RESPONSE_RETCODE_SUCCESS,allocator);
    setJsonResponse(res,json_response);
}
